/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package me.flyray.bsin.server.ws;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.annotation.PreDestroy;
import jakarta.websocket.*;
import jakarta.websocket.server.PathParam;
import me.flyray.bsin.domain.request.AppChatDTO;
import me.flyray.bsin.security.contex.LoginInfoContextHelper;
import me.flyray.bsin.security.filter.WebsocketLoginInfoInterceptor;
import me.flyray.bsin.server.biz.AppAgentEngine;
import me.flyray.bsin.server.biz.WebSocketAiChatAsyncBiz;
import org.apache.shenyu.client.spring.websocket.annotation.ShenyuServerEndpoint;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.io.IOException;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * WebSocket服务端点 - 处理AI聊天连接
 */
@Component
@ShenyuServerEndpoint(value = "/chat/{toNo}/{chatType}", configurator = WebsocketLoginInfoInterceptor.class)
public class WsServerEndpoint {

    private static final Logger LOG = LoggerFactory.getLogger(WsServerEndpoint.class);

    // 线程池配置 - 使用命名线程池便于监控
    private static final ExecutorService WEB_SOCKET_EXECUTOR =
            Executors.newFixedThreadPool(10, r -> {
                Thread t = new Thread(r, "websocket-ai-chat-" + System.currentTimeMillis());
                t.setDaemon(true); // 设置为守护线程
                return t;
            });

    // 在线连接数统计
    private static final AtomicInteger onlineCount = new AtomicInteger(0);

    // 客户端连接集合 - 线程安全
    private static final CopyOnWriteArraySet<WsServerEndpoint> wsClientMap = new CopyOnWriteArraySet<>();

    // session池 - 用户ID -> Session映射
    private static final ConcurrentHashMap<String, Session> sessionPool = new ConcurrentHashMap<>();

    // Spring应用上下文 - 解决依赖注入问题
    private static ApplicationContext applicationContext;

    // 实例属性
    private Session session;
    private String fromNo;  // 发起人用户ID
    private String toNo;    // 目标对象ID
    private String chatType; // 聊天类型

    /**
     * 设置Spring应用上下文
     */
    public static void setApplicationContext(ApplicationContext context) {
        applicationContext = context;
    }

    /**
     * WebSocket连接建立时触发
     */
    @OnOpen
    public void onOpen(@PathParam("toNo") String toNo, @PathParam("chatType") String chatType, Session session, EndpointConfig config) {
        this.session = session;
        this.toNo = toNo;
        this.chatType = chatType;

        try {
            // 从配置中获取用户信息
            UserInfo userInfo = extractUserInfo(config);
            if (userInfo == null) {
                closeWithReason(session, "用户信息获取失败");
                return;
            }

            // 检查用户认证状态
            if (!userInfo.isAuthenticated()) {
                closeWithReason(session, "用户未认证");
                return;
            }

            // 设置登录上下文
            setupLoginContext(userInfo);
            this.fromNo = userInfo.getBizRoleTypeNo();

            // 处理重复登录 - 关闭旧连接
            handleDuplicateLogin(userInfo.getBizRoleTypeNo());

            // 注册新连接
            registerConnection(userInfo.getBizRoleTypeNo(), session);

            LOG.info("新WebSocket连接建立 - SessionId: {}, FromNo: {}, ToNo: {}, 在线数: {}",
                    session.getId(), fromNo, toNo, getOnlineCount());

        } catch (Exception e) {
            LOG.error("WebSocket连接建立失败", e);
            closeWithReason(session, "连接建立失败");
        }
    }

    /**
     * WebSocket连接关闭时触发
     */
    @OnClose
    public void onClose(Session session, CloseReason closeReason) {
        try {
            // 从连接池中移除
            wsClientMap.remove(this);
            // 从session池中移除
            if (StringUtils.hasText(fromNo)) {
                sessionPool.remove(fromNo);
            }
            // 减少在线数统计
            onlineCount.decrementAndGet();
            LOG.info("WebSocket连接关闭 - SessionId: {}, FromNo: {}, 原因: {}, 剩余在线数: {}",
                    session.getId(), fromNo, closeReason.getReasonPhrase(), getOnlineCount());
        } catch (Exception e) {
            LOG.error("WebSocket连接关闭处理失败", e);
        }
    }

    /**
     * 1、接受chat内容
     * 2、获取chat用户信息
     * 3、处理聊听内容
     * 4、返回结束标识符
     */
    @OnMessage
    public void onMessage(@PathParam("toNo") String toNo, String message) throws JsonProcessingException {
        if (!StringUtils.hasText(message)) {
            LOG.warn("接收到空消息 - FromNo: {}, ToNo: {}", fromNo, toNo);
        }
        LOG.info("接收到消息 - FromNo: {}, ToNo: {}, Message: {}", fromNo, toNo, message);

        ObjectMapper mapper = new ObjectMapper();
        JsonNode messageJson = mapper.readTree(message);

        try {
            // 获取当前用户信息
            String currentFromNo = Optional.ofNullable(LoginInfoContextHelper.getLoginUser())
                    .map(user -> user.getBizRoleTypeNo())
                    .orElse(fromNo);

            // 构建聊天请求
            AppChatDTO chatRequest = AppChatDTO.builder()
                    .type(chatType)
                    .question(String.valueOf(messageJson.get("content").asText()))
                    .fromNo(currentFromNo)
                    .toNo(toNo)
                    .build();
            // 异步处理聊天消息
            processMessageAsync(chatRequest);
        } catch (Exception e) {
            LOG.error("消息处理失败 - FromNo: {}, ToNo: {}, Message: {}", fromNo, toNo, message, e);
        }
    }

    /**
     * WebSocket错误处理
     */
    @OnError
    public void onError(Session session, Throwable error) {
        LOG.error("WebSocket连接发生错误 - SessionId: {}, FromNo: {}",
                session != null ? session.getId() : "unknown", fromNo, error);
        // 尝试关闭有问题的连接
        if (session != null && session.isOpen()) {
            try {
                session.close(new CloseReason(CloseReason.CloseCodes.UNEXPECTED_CONDITION, "服务器内部错误"));
            } catch (IOException e) {
                LOG.error("关闭错误连接失败", e);
            }
        }
    }

    /**
     * 异步处理消息
     */
    private void processMessageAsync(AppChatDTO chatRequest) {
        WEB_SOCKET_EXECUTOR.submit(() -> {
            try {
                WsServerEndpoint targetEndpoint = findWebSocketEndpoint(chatRequest.getFromNo(), chatRequest.getToNo());
                if (targetEndpoint == null) {
                    LOG.warn("未找到目标WebSocket连接 - FromNo: {}, ToNo: {}",
                            chatRequest.getFromNo(), chatRequest.getToNo());
                    return;
                }
                AppAgentEngine appAgentChatBiz = applicationContext.getBean(AppAgentEngine.class);
                WebSocketAiChatAsyncBiz asyncBiz = new WebSocketAiChatAsyncBiz(chatRequest, targetEndpoint, appAgentChatBiz);
                asyncBiz.run();

            } catch (Exception e) {
                LOG.error("异步消息处理失败", e);
            }
        });
    }

    /**
     * 查找指定用户的WebSocket连接
     */
    private WsServerEndpoint findWebSocketEndpoint(String fromNo, String toNo) {
        if (wsClientMap.isEmpty() || !StringUtils.hasText(fromNo)) {
            return null;
        }
        return wsClientMap.stream()
                .filter(ws -> Objects.equals(ws.fromNo, fromNo))
                .filter(ws -> StringUtils.hasText(toNo) ? Objects.equals(ws.toNo, toNo) : true)
                .findFirst()
                .orElse(null);
    }

    /**
     * 提取用户信息
     */
    private UserInfo extractUserInfo(EndpointConfig config) {
        try {
            String bizRoleType = (String) config.getUserProperties().get("bizRoleType");
            String bizRoleTypeNo = (String) config.getUserProperties().get("bizRoleTypeNo");
            String username = (String) config.getUserProperties().get("username");
            Boolean authenticated = (Boolean) config.getUserProperties().get("authenticated");
            if (!StringUtils.hasText(bizRoleTypeNo)) {
                LOG.warn("用户角色类型编号为空");
                return null;
            }
            return new UserInfo(bizRoleType, bizRoleTypeNo, username,
                    authenticated != null ? authenticated : false);
        } catch (Exception e) {
            LOG.error("提取用户信息失败", e);
            return null;
        }
    }

    /**
     * 设置登录上下文
     */
    private void setupLoginContext(UserInfo userInfo) {
        LoginInfoContextHelper.set("bizRoleType", userInfo.getBizRoleType());
        LoginInfoContextHelper.set("bizRoleTypeNo", userInfo.getBizRoleTypeNo());
        LoginInfoContextHelper.set("username", userInfo.getUsername());
    }

    /**
     * 处理重复登录
     */
    private void handleDuplicateLogin(String bizRoleTypeNo) {
        Session existingSession = sessionPool.get(bizRoleTypeNo);
        if (existingSession != null && existingSession.isOpen()) {
            try {
                LOG.info("检测到重复登录，关闭旧连接 - 用户: {}", bizRoleTypeNo);
                sessionPool.remove(bizRoleTypeNo);
                existingSession.close(new CloseReason(CloseReason.CloseCodes.NORMAL_CLOSURE, "重复登录"));
            } catch (IOException e) {
                LOG.error("关闭重复登录连接失败 - 用户: {}", bizRoleTypeNo, e);
            }
        }
    }

    /**
     * 注册新连接
     */
    private void registerConnection(String bizRoleTypeNo, Session session) {
        sessionPool.put(bizRoleTypeNo, session);
        wsClientMap.add(this);
        onlineCount.incrementAndGet();
    }

    /**
     * 安全关闭连接
     */
    private void closeWithReason(Session session, String reason) {
        try {
            if (session != null && session.isOpen()) {
                session.close(new CloseReason(CloseReason.CloseCodes.VIOLATED_POLICY, reason));
            }
        } catch (IOException e) {
            LOG.error("关闭WebSocket连接失败: {}", reason, e);
        }
    }

    /**
     * 发送消息到客户端
     */
    public void sendMessage(String message) {
        if (session == null || !session.isOpen()) {
            LOG.warn("Session已关闭，无法发送消息 - FromNo: {}", fromNo);
            return;
        }
        try {
            synchronized (session) {
                session.getBasicRemote().sendText(message);
            }
        } catch (IOException e) {
            LOG.error("发送消息失败 - FromNo: {}, Message: {}", fromNo, message, e);
        }
    }

    public void sendMessageAsync(String message) {
        if (session.isOpen()) {
            session.getAsyncRemote().sendText(message, new SendHandler() {
                @Override
                public void onResult(SendResult result) {
                    if (!result.isOK()) {
                        // Handle send failure
                        LOG.error("Failed to send message", result.getException());
                    }
                }
            });
        }
    }

    // 静态方法 - 获取统计信息
    public static int getOnlineCount() {
        return onlineCount.get();
    }

    public static CopyOnWriteArraySet<WsServerEndpoint> getWsClientMap() {
        return new CopyOnWriteArraySet<>(wsClientMap); // 返回副本，避免外部修改
    }

    public static ConcurrentHashMap<String, Session> getSessionPool() {
        return new ConcurrentHashMap<>(sessionPool); // 返回副本，避免外部修改
    }

    // 内部用户信息类
    private static class UserInfo {
        private final String bizRoleType;
        private final String bizRoleTypeNo;
        private final String username;
        private final boolean authenticated;
        public UserInfo(String bizRoleType, String bizRoleTypeNo, String username, boolean authenticated) {
            this.bizRoleType = bizRoleType;
            this.bizRoleTypeNo = bizRoleTypeNo;
            this.username = username;
            this.authenticated = authenticated;
        }

        public String getBizRoleType() { return bizRoleType; }
        public String getBizRoleTypeNo() { return bizRoleTypeNo; }
        public String getUsername() { return username; }
        public boolean isAuthenticated() { return authenticated; }
    }

}